package com.example.toy.core.partition;

import com.example.toy.core.utils.SchemaHelper;
import lombok.SneakyThrows;
import org.apache.avro.Schema;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * HivePartitionPathStyle
 *
 * @author shenbing
 * @date 2021/12/24
 */
public class HivePartitionPathStyle implements PartitionPathStyle {

  private final Path dataDir;

  public HivePartitionPathStyle(Path dataDir) {
    this.dataDir = dataDir;
  }

  @Override
  public Path getDataDir() {
    return this.dataDir;
  }

  @Override
  public Path generate(Schema schema, Object... params) {
    if (schema == null || schema.getFields().size() == 0) {
      return this.dataDir;
    }
    if (schema.getFields().size() != params.length) {
      throw new IllegalArgumentException("partition schema fields size not equals params");
    }
    StringBuilder sb = new StringBuilder();
    for (int i = 0; i < schema.getFields().size(); i++) {
      sb.append(schema.getFields().get(i).name() + "=" + params[i]);
      sb.append("/");
    }
    String partitionPath = sb.toString();
    return new Path(this.dataDir, partitionPath);
  }

  @SneakyThrows
  @Override
  public List<Object> fromPartitionPath(Schema schema, Path path) {
    FileSystem fs = this.dataDir.getFileSystem(new Configuration());
    String dataDir = fs.makeQualified(this.dataDir).toString();
    String partitionDir = fs.makeQualified(path).toString();

    if (!partitionDir.startsWith(dataDir)) {
      throw new IllegalArgumentException("partition path is not start with data dir path");
    }
    String substring = partitionDir.substring(dataDir.length());
    String[] split = substring.split("/");

    List<Object> list = new ArrayList<>();
    int schemaIndex = 0;
    for (int i = 0; i < split.length; i++) {
      if (split[i].isEmpty()) {
        continue;
      }
      if (schemaIndex >= schema.getFields().size()) {
        throw new IllegalArgumentException("Partition path split has more fields than schema");
      }
      String value = split[i].split("=")[1];
      list.add(SchemaHelper.convert(value, schema.getFields().get(schemaIndex++).schema()));
    }
    if (list.size() != schema.getFields().size()) {
      throw new IllegalArgumentException("Partition path split fields num not match schema fields");
    }

    return list;
  }

  public static void main(String[] args) {
    Schema user = Schema.createRecord("user", "", "", false);
    user.setFields(
        Arrays.asList(
            new Schema.Field("id", Schema.create(Schema.Type.STRING)),
            new Schema.Field("name", Schema.create(Schema.Type.INT))));

    HivePartitionPathStyle hivePartitionPathStyle = new HivePartitionPathStyle(new Path("/opt"));

    Path partitionPath = hivePartitionPathStyle.generate(user, "aa", 1);
    System.out.println(partitionPath);

    List<Object> objects =
        hivePartitionPathStyle.fromPartitionPath(user, new Path("file:///opt/a=1/c=112//"));
    System.out.println(objects);
  }
}
